# coding=utf-8
from torchvision.models import resnet
from torchvision.models import inception
import torch.nn as nn


def choose(arch, num_classes, pretrained=False):
    if(arch == 'resnet50'):
        model = resnet.__dict__[arch](pretrained)
        if(pretrained):
            for param in model.parameters():  # nn.Module有成员函数parameters()
                param.requires_grad = False
        model.fc = nn.Sequential(nn.Linear(512 * resnet.Bottleneck.expansion, num_classes), nn.Sigmoid())
        return model
    elif(arch == 'resnet18'):
        model = resnet.__dict__[arch](pretrained)
        if(pretrained):
            for param in model.parameters():  # nn.Module有成员函数parameters()
                param.requires_grad = False
        model.fc = nn.Sequential(nn.Linear(512 * resnet.BasicBlock.expansion, num_classes), nn.Sigmoid())
        return model
    elif(arch == 'inception_v3'):
        model = inception.__dict__[arch](pretrained)
        if(pretrained):
            for param in model.parameters():  # nn.Module有成员函数parameters()
                param.requires_grad = False
        model.fc = nn.Sequential(nn.Linear(2048, num_classes), nn.Sigmoid())
        model.aux_logits = False
        return model
